import cv2 as cv
import numpy as np
import os
import pickle
from pic_utils import pic_tools,txt_tools
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import MinMaxScaler
class insects_tools():

    @staticmethod
    def count_insects(img):
        Prewitt = pic_tools.prewitt_operator(img) #prewitt算子进行边缘检测
        contours = pic_tools.getContours(Prewitt)
        count=0
        for cont in contours:
            ares = cv.contourArea(cont)
            if ares < 5 or ares > 120:
                continue
            count += 1
        return count

    @staticmethod
    def insects_count_visualization(img_path):
        img = pic_tools.read_imgs(img_path)
        Prewitt=pic_tools.prewitt_operator(img)
        contours = pic_tools.getContours(Prewitt)
        count = 0
        for cont in contours:
            ares = cv.contourArea(cont)
            if ares < 5 or ares > 120:
                continue
            count += 1
            print('标号{}--面积{}:'.format(count, ares), end='')
            x, y, w, h = cv.boundingRect(cont)
            print('坐标为x-{}-y-{}'.format(x, y))
            cv.rectangle(img, (x, y, w, h), (0, 0, 255), 1)  # 为每只昆虫绘制外接矩形
            y = 10 if y < 10 else y
            cv.putText(img, str(count), (x, y), cv.FONT_HERSHEY_COMPLEX, 0.4, (0, 255, 0), 1)
        cv.namedWindow('count_results', 2)
        cv.imshow('count_results', img)
        cv.waitKey(0)

    # 统计一张图片上的虫群数量
    @staticmethod
    def count_insects_on_single_img(img_path):
        img = pic_tools.read_imgs(img_path)
        # processed_img = pic_tools.cutout_imgs(img)  # 截取粘虫板
        insects_num = insects_tools.count_insects(img)  # 对粘虫板进行计数
        print(str(img_path).rsplit('\\',1)[1],insects_num)
        return insects_num

    # 统计一张粘虫板上的虫群数量
    @staticmethod
    def count_insects_on_single_board(a_side, b_side):
        _, filename = os.path.split(a_side)  # 获取图片名称
        board_location = str(filename).split(".")[0][0:5]
        insects_board = insects_tools.count_insects_on_single_img(a_side) + insects_tools.count_insects_on_single_img(b_side)
        return board_location, insects_board

    # 统计一个温室每张粘虫板上的虫群数量
    # input: 一个温室的图片路径
    # output(dict): 该温室下每张粘虫板的虫群数量以及虫群总数
    @staticmethod
    def green_house_insects_info(greenhouse_gallery_path):
        imgs_dir = pic_tools.read_imgs_from_dir(greenhouse_gallery_path)
        gh_insects_info = {}
        gh_insects_boards_info={}
        insects_total = 0
        img_index = list(range(0, len(imgs_dir)))
        greenhouse_tag = str(imgs_dir[0]).rsplit('\\',1)[1][0]#选出这是几号大棚
        for a_side, b_side in zip(img_index[::2], img_index[1::2]):  # 根据图片遍历的奇偶匹配一张粘虫板的a，b面
            board_location, insects_num_board = insects_tools.count_insects_on_single_board(imgs_dir[a_side], imgs_dir[b_side])
            gh_insects_boards_info[board_location] = insects_num_board
        for insects_num in gh_insects_boards_info.values():  # 获取温室所有粘虫板上害虫的总和
            insects_total += insects_num
        gh_insects_info["each_board_info"]=gh_insects_boards_info
        gh_insects_info["insects_total"] = insects_total
        gh_insects_info["greenhouse_tag"] = greenhouse_tag
        return gh_insects_info

    # 为单张虫害图像区分A，B昆虫并分别计数
    @staticmethod
    def classify_insects_on_single_chas_file(chas_file_path,model_path,scaler):
        X_test = np.loadtxt(chas_file_path)
        X_test_scaler=scaler.transform(X_test)
        with open(model_path, 'rb') as model_file:
            model = pickle.load(model_file)
        result=model.predict(X_test_scaler)
        insect_A=sum(result==0)
        insect_B=sum(result==1)
        insect_C=sum(result==2)
        return insect_A,insect_B,insect_C

    # 为单张粘虫板区分A，B昆虫并分别计数
    @staticmethod
    def classify_insects_on_single_board(a_side,b_side,model_path,scaler):
        _, filename = os.path.split(a_side)  # 获取图片名称
        board_location = str(filename).split(".")[0][0:5]
        insect_A_from_a,insect_B_from_a,insect_C_from_a = insects_tools.classify_insects_on_single_chas_file(a_side,model_path,scaler)
        insect_A_from_b,insect_B_from_b,insect_C_from_b = insects_tools.classify_insects_on_single_chas_file(b_side, model_path,scaler)
        insect_A=insect_A_from_a+insect_A_from_b
        insect_B=insect_B_from_a+insect_B_from_b
        return board_location, insect_A,insect_B

    @staticmethod
    def classify_green_house_insects(chas_file_dir,model,scaler_path):
        chas_paths=txt_tools.read_files_from_dir(chas_file_dir)
        #读取对应的数据缩放器
        with open(scaler_path, 'rb') as scaler_file:
            scaler = pickle.load(scaler_file)
        gh_insects_info = {}
        gh_insects_boards_info={}
        insects_total = 0
        chas_index = list(range(0, len(chas_paths)))
        greenhouse_tag = str(chas_paths[0]).rsplit('\\',1)[1][0]#把对应的大棚数量提取出来。
        insect_A_total=0
        insect_B_total=0
        for a_side, b_side in zip(chas_index[::2], chas_index[1::2]):  # 根据图片遍历的奇偶匹配一张粘虫板的a，b面
            board_location, insect_A,insect_B = insects_tools.classify_insects_on_single_board(chas_paths[a_side], chas_paths[b_side],model,scaler)
            gh_insects_boards_info[board_location] = {'飞虱':insect_A,'蓟马':insect_B}
            insect_A_total+=insect_A
            insect_B_total+=insect_B
        gh_insects_info["each_board_info"]=gh_insects_boards_info
        gh_insects_info["insects_total_each"] = {'飞虱总数':insect_A_total,'蓟马总数':insect_B_total}
        gh_insects_info["insects_total"] = insect_A_total+insect_B_total
        gh_insects_info["greenhouse_tag"] = greenhouse_tag
        return gh_insects_info

